import cv2
import numpy as np
import os
import time
import threading
from concurrent.futures import ThreadPoolExecutor
import atexit
import datetime
from collections import deque

class MergedPosePositionVisualizer:
    """预计算优化的合并可视化器 - 标准pose和vector在初始化时设定"""

    def __init__(self, st_pose=None, st_vector=None, rect_size=(0.8, 2.5), 
                 scale_factor=100, enable_display=True):
        """
        初始化预计算可视化器
        
        Args:
            st_pose_dict: 标准pose字典，初始化时设定
            standard_vector: 标准位置向量，初始化时设定
            rect_size: 矩形大小（宽, 长），单位为米，默认为(0.8, 2.5)
            scale_factor: 米到像素的转换因子，默认为100（1米=100像素）
            enable_display: 是否启用显示功能，默认为True
        """
        # 存储标准数据
        self.st_pose_dict = st_pose
        self.standard_vector = st_vector
        
        # Position部分的参数
        self.rect_size = rect_size
        self.scale_factor = scale_factor
        self.rect_width, self.rect_height = rect_size
        
        # 画布配置 - 使用固定尺寸避免缩放
        self.scene_width = 1280  # 左边场景区域：固定1280x720
        self.scene_height = 720
        self.position_width = 400  # 右边位置区域：400像素宽
        self.canvas_width = self.scene_width + self.position_width  # 总宽度1680
        self.canvas_height = self.scene_height  # 总高度720
        
        # 显示控制
        self.enable_display = enable_display
        self.window_created = False
        self.window_name = "Optimized Merged Visualization"
        
        # 预分配画布 - 关键优化
        self.canvas = np.zeros((self.canvas_height, self.canvas_width, 3), dtype=np.uint8)
        self.scene_canvas = self.canvas[:, :self.scene_width]  # 1280x720区域
        self.position_canvas = self.canvas[:, self.scene_width:]  # 400x720区域
        
        # 位置图的中心点和比例
        self.pos_center_x = self.position_width // 2
        self.pos_center_y = self.scene_height // 2
        self.pos_scale = min(self.position_width, self.scene_height) // 4
        
        # Pose部分的颜色定义
        self.pose_colors = {
            'left_top': (0, 255, 0),    # 绿色
            'left_0': (255, 0, 255),    # 紫色
            'left_bottom': (255, 0, 0),  # 蓝色
            'bottom_0': (0, 128, 128),   # 深青色
            'right_bottom': (0, 255, 255), # 黄色
            'right_0': (255, 255, 0),    # 青色
            'right_top': (0, 0, 255),    # 红色
            'top_0': (128, 0, 128),      # 深紫色
        }
        
        # 整体颜色配置
        self.colors = {
            'standard_pose': (0, 255, 0),      # 绿色
            'current_pose': (0, 0, 255),       # 红色  
            'standard_rect': (0, 255, 255),    # 黄色
            'current_rect': (255, 0, 255),     # 紫色
            'background': (40, 40, 40),        # 深灰背景
            'grid': (80, 80, 80),              # 网格线
            'text': (255, 255, 255),           # 白色文字
            'split_line': (128, 128, 128)      # 分割线
        }
        
        # 定义连接顺序
        self.connect_order = [
            'left_top', 'left_0', 'left_bottom',
            'bottom_0', 'right_bottom', 'right_0',
            'right_top', 'top_0', 'left_top'
        ]
        
        # 字体设置
        self.font = cv2.FONT_HERSHEY_SIMPLEX
        self.font_scale = 0.4
        self.font_thickness = 1
        
        # 输出目录
        self.output_dir = "results_optimized_visualization"
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 性能监控
        self.frame_times = deque(maxlen=30)
        self.fps = 0
        
        # 图像缓存优化
        self._image_cache = {}
        self.cache_max_size = 3  # 减少缓存大小
        
        # === 预计算优化部分 ===
        # 预计算标准pose的点和连线
        self._precompute_standard_elements()
        
        # 预绘制静态元素（包含标准pose和标准vector）
        self._prepare_static_background()
        
        # 注册清理函数
        atexit.register(self.cleanup)
        
        print(f"预计算可视化器初始化完成，输出目录: {self.output_dir}")

    def _precompute_standard_elements(self):
        """预计算标准pose和vector的几何信息"""
        # 预计算标准pose的有效点
        self.st_pose_valid_points = []
        if self.st_pose_dict:
            for name in self.connect_order[:-1]:  # 去掉重复的最后一个点
                if name in self.st_pose_dict and self.st_pose_dict[name]:
                    x, y = self.st_pose_dict[name]
                    point = (int(x), int(y))
                    self.st_pose_valid_points.append(point)
        
        # 预计算标准vector的矩形点
        self.standard_rect_points = None
        if self.standard_vector:
            self.standard_rect_points = self._compute_rectangle_points(self.standard_vector)

    def _prepare_static_background(self):
        """预绘制包含标准元素的静态背景"""
        # 清空画布为深色背景
        self.canvas[:] = self.colors['background']
        
        # 绘制分割线
        cv2.line(self.canvas, (self.scene_width, 0), (self.scene_width, self.scene_height), 
                self.colors['split_line'], 2)
        
        # 绘制标题
        cv2.putText(self.canvas, "POSE (Optimized)", (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                   1, self.colors['text'], 2)
        cv2.putText(self.canvas, "POSITION", (self.scene_width + 20, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, self.colors['text'], 2)
        
        # 在位置区域绘制网格
        self._draw_position_grid()
        
        # === 预绘制标准pose ===
        if self.st_pose_dict and self.st_pose_valid_points:
            # 绘制标准pose的点
            for name in self.connect_order[:-1]:
                if name in self.st_pose_dict and self.st_pose_dict[name]:
                    x, y = self.st_pose_dict[name]
                    point = (int(x), int(y))
                    cv2.circle(self.scene_canvas, point, 3, (192, 192, 192), -1)
            
            # 绘制标准pose的连线
            if len(self.st_pose_valid_points) >= 3:
                cv2.polylines(self.scene_canvas, [np.array(self.st_pose_valid_points, dtype=np.int32)], 
                             True, self.colors['standard_pose'], 2)
        
        # === 预绘制标准vector ===
        if self.standard_rect_points is not None:
            cv2.polylines(self.position_canvas, [self.standard_rect_points], True, 
                         self.colors['standard_rect'], 2)
        
        # 保存静态背景（包含标准元素）
        self.static_background_with_standards = self.canvas.copy()

    def _draw_position_grid(self):
        """绘制位置区域的网格"""
        # 绘制十字参考线
        cv2.line(self.position_canvas, 
                (self.pos_center_x, 0), (self.pos_center_x, self.scene_height),
                self.colors['grid'], 1)
        cv2.line(self.position_canvas,
                (0, self.pos_center_y), (self.position_width, self.pos_center_y), 
                self.colors['grid'], 1)
        
        # 绘制同心圆作为距离参考
        for radius in [50, 100, 150]:
            if radius < min(self.pos_center_x, self.pos_center_y):
                cv2.circle(self.position_canvas, 
                          (self.pos_center_x, self.pos_center_y), 
                          radius, self.colors['grid'], 1)

    def _create_display_window(self):
        """创建显示窗口（只创建一次）"""
        if not self.enable_display or self.window_created:
            return
            
        try:
            cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
            cv2.resizeWindow(self.window_name, self.canvas_width, self.canvas_height)
            self.window_created = True
        except Exception as e:
            print(f"无法创建显示窗口: {e}")
            self.enable_display = False

    def _process_background_image(self, image_input):
        """超高效处理背景图像"""
        if image_input is None:
            return None
            
        # 处理图像输入
        if isinstance(image_input, (str, os.PathLike)):
            image_path = str(image_input)
            
            # 检查缓存
            if image_path in self._image_cache:
                image = self._image_cache[image_path]
            else:
                image = cv2.imread(image_path)
                if image is None:
                    return None
                
                # 有限缓存
                if len(self._image_cache) < self.cache_max_size:
                    self._image_cache[image_path] = image
                    
        elif isinstance(image_input, np.ndarray):
            image = image_input
        else:
            return None
        
        # 检查尺寸是否匹配
        h, w = image.shape[:2]
        if h == self.scene_height and w == self.scene_width:
            # 完美匹配，直接返回
            return image
        elif h == self.scene_height and w != self.scene_width:
            # 高度匹配，宽度裁剪或填充
            if w > self.scene_width:
                # 裁剪中间部分
                start_x = (w - self.scene_width) // 2
                return image[:, start_x:start_x + self.scene_width]
            else:
                # 宽度不足，居中填充
                result = np.zeros((self.scene_height, self.scene_width, 3), dtype=np.uint8)
                start_x = (self.scene_width - w) // 2
                result[:, start_x:start_x + w] = image
                return result
        else:
            # 尺寸不匹配，只在必要时才resize
            return cv2.resize(image, (self.scene_width, self.scene_height))

    def _draw_current_pose(self, pose_dict):
        """高效绘制当前pose点和连线"""
        if pose_dict is None:
            return
        
        # 收集有效点
        valid_points = []
        
        for name in self.connect_order[:-1]:  # 去掉重复的最后一个点
            if name in pose_dict and pose_dict[name]:
                x, y = pose_dict[name]
                point = (int(x), int(y))
                valid_points.append(point)
                
                # 绘制点 - 使用单独颜色
                color = self.pose_colors.get(name, self.colors['current_pose'])
                cv2.circle(self.scene_canvas, point, 5, color, -1)
                
                # 添加标签（简化版，只有关键点）
                if name in ['left_top', 'right_top', 'left_bottom', 'right_bottom']:
                    cv2.putText(self.scene_canvas, name.replace('_', ''), 
                              (point[0]+8, point[1]), self.font, 
                              self.font_scale, color, self.font_thickness)
        
        # 绘制连线（闭合多边形）
        if len(valid_points) >= 3:
            cv2.polylines(self.scene_canvas, [np.array(valid_points, dtype=np.int32)], 
                         True, self.colors['current_pose'], 3)

    def _compute_rectangle_points(self, vector):
        """快速计算矩形四个角点"""
        if vector is None or 'center' not in vector:
            return None
            
        center_x = vector['center'][0] * self.scale_factor + self.pos_center_x
        center_y = vector['center'][1] * self.scale_factor + self.pos_center_y
        
        # 默认矩形大小（像素）
        rect_w = self.rect_size[0] * self.scale_factor
        rect_h = self.rect_size[1] * self.scale_factor
        
        # 简化处理方向向量
        if 'left_vec' in vector and 'bottom_vec' in vector:
            left_vec = np.array(vector['left_vec'][:2])
            bottom_vec = np.array(vector['bottom_vec'][:2])
            
            # 快速标准化
            left_len = np.linalg.norm(left_vec)
            if left_len > 1e-6:
                left_vec = left_vec / left_len * rect_h
            else:
                left_vec = np.array([0, rect_h])
                
            bottom_len = np.linalg.norm(bottom_vec)
            if bottom_len > 1e-6:
                bottom_vec = bottom_vec / bottom_len * rect_w
            else:
                bottom_vec = np.array([rect_w, 0])
        else:
            # 默认方向
            left_vec = np.array([0, rect_h])
            bottom_vec = np.array([rect_w, 0])
        
        # 计算四个角点
        center = np.array([center_x, center_y])
        bottom_left = center
        bottom_right = center + bottom_vec
        top_left = center + left_vec  
        top_right = top_left + bottom_vec
        
        return np.array([bottom_left, bottom_right, top_right, top_left], dtype=np.int32)

    def _draw_current_rectangle(self, vector):
        """在位置区域绘制当前矩形"""
        if vector is None:
            return
            
        points = self._compute_rectangle_points(vector)
        if points is not None:
            cv2.polylines(self.position_canvas, [points], True, 
                         self.colors['current_rect'], 3)

    def merge_visualize(self, pose, vector, color_img, 
                        timestamp=None, display=True, save_async=True, 
                        fps=0.0, frame_count=0):
            """
            优化的合并可视化函数 - 只需传入当前数据
            接收外部计算的FPS和帧数信息
            
            Args:
                pose: 当前pose字典
                vector: 当前测试位置向量
                color_img: 彩色图像 (numpy.ndarray)
                timestamp: 时间戳，用作文件名(可选)
                display: 是否显示图像（默认为True）
                save_async: 是否异步保存（默认为True）
                fps: 外部计算的系统FPS
                frame_count: 当前帧数
            
            Returns:
                tuple: (output_path, canvas_copy)
            """
            try:
                # 1. 快速恢复包含标准元素的静态背景
                self.canvas[:] = self.static_background_with_standards
                
                # 2. 高效处理和绘制背景图像
                processed_image = self._process_background_image(color_img)
                if processed_image is not None:
                    # 创建带透明度的混合效果（可选）
                    alpha = 0.7  # 背景图透明度
                    cv2.addWeighted(self.scene_canvas, 1-alpha, processed_image, alpha, 0, self.scene_canvas)
                
                # 3. 绘制当前pose（标准pose已在静态背景中）
                if pose:
                    self._draw_current_pose(pose)
                
                # 4. 绘制当前位置矩形（标准矩形已在静态背景中）
                if vector:
                    self._draw_current_rectangle(vector)
                
                # 5. 添加检测状态信息
                status_text = "Detection: "
                if pose and vector:
                    status_text += "ACTIVE"
                    status_color = (0, 255, 0)  # 绿色
                else:
                    status_text += "NO RESULT"
                    status_color = (0, 0, 255)  # 红色
                
                cv2.putText(self.canvas, status_text, (20, self.scene_height - 50), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, status_color, 2)
                
                # 6. 显示系统FPS和帧数信息
                fps_text = f"System FPS: {fps:.1f} | Frame: {frame_count}"
                cv2.putText(self.canvas, fps_text, (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, (0, 255, 255), 2)  # 黄色文字，更显眼
                
                # 7. 显示图像
                if display and self.enable_display:
                    self._create_display_window()
                    if self.window_created:
                        try:
                            cv2.imshow(self.window_name, self.canvas)
                            key = cv2.waitKey(1) & 0xFF
                            if key == ord('q'):
                                print("用户请求退出")
                                return None, None  # 可以在这里处理退出信号
                        except Exception as e:
                            print(f"显示图像失败: {e}")
                
                # 8. 保存图像（如果需要）
                output_path = None
                if save_async and timestamp:
                    filename = f"optimized_{timestamp}.png"
                    output_path = os.path.join(self.output_dir, filename)
                    
                    # 快速保存
                    cv2.imwrite(output_path, self.canvas, [cv2.IMWRITE_JPEG_QUALITY, 70])
                
                return output_path, self.canvas.copy()
                
            except Exception as e:
                print(f"可视化过程中发生错误: {e}")
                return None, None
            
    def update_standards(self, new_st_pose_dict=None, new_standard_vector=None):
        """运行时更新标准pose和vector"""
        if new_st_pose_dict is not None:
            self.st_pose_dict = new_st_pose_dict
        if new_standard_vector is not None:
            self.standard_vector = new_standard_vector
        
        # 重新预计算
        self._precompute_standard_elements()
        self._prepare_static_background()
        print("标准元素已更新并重新预计算")

    def clear_cache(self):
        """清空缓存"""
        self._image_cache.clear()

    def cleanup(self):
        """清理资源"""
        try:
            if self.window_created:
                cv2.destroyWindow(self.window_name)
            
            self.clear_cache()
        except Exception as e:
            print(f"清理资源时出错: {e}")

    def __del__(self):
        """析构函数"""
        self.cleanup()


def demo_optimized_visualizer():
    """演示优化后的可视化器"""
    print("=== 预计算优化可视化器演示 ===")
    
    # 定义标准pose和vector（初始化时设定）
    st_pose_dict = {
        'left_top': (150, 100), 'left_0': (150, 200), 
        'left_bottom': (150, 300), 'bottom_0': (250, 300),
        'right_bottom': (350, 300), 'right_0': (350, 200),
        'right_top': (350, 100), 'top_0': (250, 100)
    }
    
    standard_vector = {
        'center': [0, 0],
        'left_vec': [0, 2.5, 0],
        'bottom_vec': [0.8, 0, 0]
    }
    
    # 创建优化的可视化器（标准数据在初始化时传入）
    visualizer = MergedPosePositionVisualizer(
        st_pose_dict=st_pose_dict,
        standard_vector=standard_vector,
        enable_display=True
    )
    
    print("开始实时演示... (按 'q' 退出)")
    start_time = time.time()
    frame_count = 0
    
    try:
        while True:
            # 生成当前数据
            t = time.time() - start_time
            
            # 生成背景图像（固定1280x720）
            bg = np.random.randint(50, 200, (720, 1280, 3), dtype=np.uint8)
            
            # 动画当前pose
            offset_x = int(30 * np.sin(t * 2))
            offset_y = int(20 * np.cos(t * 1.5))
            
            current_pose = {}
            for key, (x, y) in st_pose_dict.items():
                current_pose[key] = (x + offset_x, y + offset_y)
            
            # 动画当前vector
            current_vector = {
                'center': [0.5 * np.sin(t), 0.3 * np.cos(t * 1.2)],
                'left_vec': [np.sin(t * 0.5) * 2.5, np.cos(t * 0.5) * 2.5, 0],
                'bottom_vec': [np.cos(t * 0.3) * 0.8, np.sin(t * 0.3) * 0.8, 0]
            }
            
            # === 关键：只需传入当前数据 ===
            result, frame = visualizer.merge_visualize(
                pose_dict=current_pose,
                test_vector=current_vector,
                color_img=bg,
                timestamp=f"demo_{frame_count:04d}",
                display=True,
                save_async=False
            )
            
            frame_count += 1
            
            # 检查退出
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
                
            # 限制演示时间
            if t > 20:  # 20秒后自动结束
                print("演示时间结束")
                break
                
    except KeyboardInterrupt:
        print("\n用户中断演示")
    
    elapsed = time.time() - start_time
    avg_fps = frame_count / elapsed if elapsed > 0 else 0
    
    print(f"\n演示统计:")
    print(f"总时间: {elapsed:.1f}s")
    print(f"总帧数: {frame_count}")
    print(f"平均FPS: {avg_fps:.1f}")
    
    visualizer.cleanup()


if __name__ == "__main__":
    demo_optimized_visualizer()